# Read vtk file produced by Snoopy
# 
#
# Usage: 
#   import readVTK as r; 
#   V = r.readVTK(vtkfile)
#
#   V:       Structure in which the data are stored
#   vtkfile: The filename
#   notes:   Only reads binary STRUCTURED_POINTS
#
# 
# Based on MATLAB scripts by
# Erik Vidholm 2006
# Geoffroy Lesur 2009
#       Extended to include several fields in the same file (as Snoopy does)
#       The output is now a structure.
#
# Translated to python by Jeremy Goodman 2013
#
# The output is a dictionary; e.g., V['vx'] is an ndarray vx[:,:,:]
#
from numpy import *

def readVTK(vtkfile):

    try:
        fid = open(vtkfile,'rb')
    except IOError as e:
        print "I/O error({0}): {1}".format(e.errno, e.strerror)
        raise

    s = fid.readline() #  "vtk DataFile Version x.x"
    s = fid.readline() #  "comments"
    s = fid.readline() #  "BINARY"
    s = fid.readline() #  "DATASET STRUCTURED_POINTS"

    s = fid.readline() #  "DIMENSIONS NX NY NZ"
    sz = [int(n) for n in s.split()[1:4]]

    s = fid.readline() #  "ORIGIN OX OY OZ"
    origin = [double(x) for x in s.split()[1:4]]

    s = fid.readline() #  "SPACING OX OY OZ"
    spacing = [double(x) for x in s.split()[1:4]]    

    # Generate coordinates
    x = origin[0] + spacing[0]*arange(sz[0])
    y = origin[1] + spacing[1]*arange(sz[1])
    z = origin[2] + spacing[2]*arange(sz[2])

    # Initialize the output dictionary
    V = {'x': x, 'y': y, 'z': z, 'nx': sz[0], 'ny': sz[1], 'nz': sz[2]}

    s = fid.readline() # "POINT_DATA NXNYNZ"
    s = fid.readline() # "SCALARS/VECTORS name data_type" (e.g: "SCALARS vx float")
    varname = s.split()[1]

    # The first one is treated as a scalar
    s = fid.readline() # "LOOKUP_TABLE default"
    dt = dtype('>f')
    V[varname] = fromfile(fid,dtype=dt,count=prod(sz)).reshape(sz)

    # Get the other components from the FIELD
    s = fid.readline() # "FIELD Fieldname num_field"
    num_field = int(s.split()[2])

    #Loop over the remaining fields
    for i in range(num_field):
        s = fid.readline() # fieldname dimensionality num_point data_type
        varname = s.split()[0]
        V[varname] = fromfile(fid,dtype=dt,count=prod(sz)).reshape(sz)

    fid.close()
    return V
